#!/usr/bin/env python3
"""
简单测试版本 - 先测试基本功能
"""

import cv2
import time
import threading
from PIL import Image
import xgoscreen.LCD_2inch as LCD_2inch
from ultralytics import YOLO
import os

class SimpleTest():
    def __init__(self):
        # 初始化显示屏
        self.display = LCD_2inch.LCD_2inch()
        self.display.Init() 
        self.display.clear()
        
        # 摄像头相关
        self.cap = None
        self.camera_still = False
        
        # 加载YOLO模型
        if os.path.exists("best.pt"):
            self.model = YOLO("best.pt")
            print("✅ YOLO模型加载成功")
        else:
            self.model = None
            print("⚠️ 未找到YOLO模型文件 best.pt")
    
    def open_camera(self):
        """打开摄像头"""
        print("📷 正在打开摄像头...")
        self.cap = cv2.VideoCapture(0)
        
        if not self.cap.isOpened():
            print("❌ 无法打开摄像头")
            return False
        
        # 设置摄像头参数
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 240)
        self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
        self.cap.set(cv2.CAP_PROP_FPS, 30)
        
        print("✅ 摄像头打开成功")
        return True
    
    def test_basic_camera(self):
        """测试基本摄像头显示"""
        print("🎥 测试基本摄像头显示...")
        
        if not self.open_camera():
            return
        
        self.camera_still = True
        frame_count = 0
        
        while self.camera_still:
            ret, frame = self.cap.read()
            if not ret:
                print("无法读取摄像头图像")
                time.sleep(0.1)
                continue
            
            frame_count += 1
            
            # 在图像上添加测试信息
            cv2.putText(frame, f"Frame: {frame_count}", (10, 30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
            cv2.putText(frame, "Basic Camera Test", (10, 60), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
            cv2.putText(frame, "Press Ctrl+C to exit", (10, 90), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2)
            
            # 显示图像
            try:
                image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                image_rgb = cv2.flip(image_rgb, 1)  # 水平翻转
                pil_image = Image.fromarray(image_rgb)
                self.display.ShowImage(pil_image)
            except Exception as e:
                print(f"显示图像错误: {e}")
            
            time.sleep(0.05)
    
    def test_yolo_detection(self):
        """测试YOLO检测"""
        print("🧠 测试YOLO检测...")
        
        if not self.open_camera():
            return
        
        if self.model is None:
            print("❌ YOLO模型未加载")
            return
        
        self.camera_still = True
        frame_count = 0
        
        while self.camera_still:
            ret, frame = self.cap.read()
            if not ret:
                print("无法读取摄像头图像")
                time.sleep(0.1)
                continue
            
            frame_count += 1
            
            # YOLO检测
            results = self.model.predict(frame, conf=0.3, verbose=False, imgsz=320)
            
            # 处理检测结果
            detection_count = 0
            for r in results:
                if r.boxes is not None and len(r.boxes) > 0:
                    for box in r.boxes:
                        conf = float(box.conf[0])
                        cls_id = int(box.cls[0])
                        label = self.model.names[cls_id]
                        
                        x1, y1, x2, y2 = map(int, box.xyxy[0])
                        
                        # 绘制检测框
                        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                        
                        # 绘制标签
                        text = f"{label} {conf:.2f}"
                        cv2.putText(frame, text, (x1, y1 - 10), 
                                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
                        
                        detection_count += 1
            
            # 添加帧计数和检测信息
            cv2.putText(frame, f"Frame: {frame_count}", (10, 30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
            cv2.putText(frame, f"Detections: {detection_count}", (10, 60), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
            cv2.putText(frame, "YOLO Test - Press Ctrl+C", (10, 90), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2)
            
            # 显示图像
            try:
                image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                image_rgb = cv2.flip(image_rgb, 1)
                pil_image = Image.fromarray(image_rgb)
                self.display.ShowImage(pil_image)
            except Exception as e:
                print(f"显示图像错误: {e}")
            
            time.sleep(0.05)
    
    def test_display(self):
        """测试显示屏"""
        print("🖥️ 测试显示屏...")
        
        # 显示不同颜色的测试图像
        colors = ["red", "green", "blue", "yellow", "white", "black"]
        
        for color in colors:
            print(f"显示 {color} 颜色")
            test_image = Image.new("RGB", (320, 240), color)
            self.display.ShowImage(test_image)
            time.sleep(1)
        
        # 显示黑色
        black_image = Image.new("RGB", (320, 240), "black")
        self.display.ShowImage(black_image)
        print("✅ 显示屏测试完成")
    
    def close(self):
        """清理资源"""
        self.camera_still = False
        time.sleep(0.2)
        if self.cap:
            self.cap.release()
        self.display.clear()
        print("💤 测试程序已关闭")

def main():
    print("🔧 简单测试程序")
    print("=" * 40)
    
    tester = SimpleTest()
    
    try:
        print("请选择测试项目:")
        print("1. 测试显示屏")
        print("2. 测试基本摄像头显示")
        print("3. 测试YOLO检测")
        
        choice = input("请输入选择 (1-3): ").strip()
        
        if choice == "1":
            tester.test_display()
        elif choice == "2":
            print("📷 基本摄像头测试，按 Ctrl+C 退出")
            tester.test_basic_camera()
        elif choice == "3":
            print("🧠 YOLO检测测试，按 Ctrl+C 退出")
            tester.test_yolo_detection()
        else:
            print("无效选择")
            
    except KeyboardInterrupt:
        print("\n⚠️ 用户中断程序")
    finally:
        tester.close()
        print("🎯 测试程序结束")

if __name__ == "__main__":
    main() 